# use sparse4d-v3: https://github.com/HorizonRobotics/Sparse4D
# commit id: c41df4bbf7bc82490f11ff55173abfcb3fb91425

from ..module import BaseModule
import torch.nn as nn

class NVTXSparse4Dv3(BaseModule):
    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self.depth_branch_hook = None

    def _get_submodules(self):
        self.backbone = getattr(self.module, 'img_backbone', None)
        assert isinstance(self.backbone, nn.Module), 'backbone is not found'

        self.neck = getattr(self.module, 'img_neck', None)
        assert isinstance(self.neck, nn.Module), 'neck is not found'

        self.head = getattr(self.module, 'head', None)
        assert self.head is not None, 'head is not found'

        self.depth_branch = getattr(self.module, 'depth_branch', None)
        assert self.depth_branch is not None, 'depth_branch is not found'
        self.submodule_hook_maps[self.depth_branch] = self.depth_branch_hook

        self.submodules_conformed = True